import sys, re
import elementtree.ElementTree as ET
from heapq import heappush, heappop
from collections import deque


#-------------------------------------------------------------------------------
# Penalty for missing or inserting a segment (for computing optimal alignment)
#-------------------------------------------------------------------------------

skip = 800

#-------------------------------------------------------------------------------
# Main program loop
#-------------------------------------------------------------------------------

def main():
    # We need a separate logfile for each segment
    Logfile = sys.argv[1] 
    spanList = stripLog(Logfile)
    print '\nLog "%s" contains %d words.' % (Logfile, len(spanList))
    DATA = sys.argv[2:]
    print 'Annotation files:'
    for d in DATA: print '   ',d
    #---------------------------------------------------
    # Get a base XML tree from the first annotation file
    #---------------------------------------------------
    # chain = (Name, begin, end, seg)
    oldAnnot = []
    Annoz = [extractAnnoz(dat, oldAnnot) for dat in DATA]
    starttime = sum([D[0]-2000 for D in Annoz])/len(Annoz)
    Annoz = [[starttime]+D for D in Annoz]
    chain = bestalign(Annoz)
    nuTiers = makeTier(chain,spanList)
    buildEAF(nuTiers+oldAnnot)

  
#-------------------------------------------------------------------------------
# Get the words from the log file
#-------------------------------------------------------------------------------

def stripLog(logFile):
    offset = 18     
    # offset is a number of seconds to correct the time
    # in the log file so that it matches the viedo time. 
    spanList = []
    word, T0 = 'None', None
    file = open(logFile, 'r')
    for line in file.readlines():
        T,B,W = line.split('#')
        if T[0] == '%': continue
        if not T0: 
            [H0,M0,S0] = T.split(' ')[1].split(':')[:3]
            T0 = ((((int(H0) * 60) + int(M0)) * 60) + int(S0)) + offset
        [H,M,S] = T.split(' ')[1].split(':')[:3]
        Time = ((((int(H) * 60) + int(M)) * 60) + int(S)) - T0
        button = B.strip()        
        wd = W.strip()
        if button == 'SHOW': 
            spanList.append([Time])
        elif button == 'OK':
            spanList[-1].append(wd)
            spanList[-1].append(Time)
        elif button == 'FAIL':
            spanList[-1].append('('+wd+')')
            spanList[-1].append(Time)
    return spanList
        

#-------------------------------------------------------------------------------
# Change times from milliseconds to 'seconds.milli' with three trailing digits 
#-------------------------------------------------------------------------------

def timeDec(time):
    if int(time) < 10: return '0.00' + str(time)
    if int(time) < 100: return '0.0' + str(time)
    if int(time) < 1000: return '0.' + str(time)
    return str(time)[:-3] + '.' + str(time)[-3:]


#-------------------------------------------------------------------------------
# Extract annotations from a .eaf file (which is an XML file)
#-------------------------------------------------------------------------------

def extractAnnoz(d, oldAnnot):
    tree = ET.parse(d)
    root = tree.getroot()
    if len(root) > 9: 
        print 'Expected 1 tier of annotations,'
        print 'only extractng the first one in "%s".' % d
    pairs = [(X.attrib['TIME_SLOT_ID'],X.attrib['TIME_VALUE']) for X in root[1]]
    timeDict = dict(pairs)
    tName = root[2].attrib['TIER_ID']
    Annoz = []
    for xx in root[2]:
        time1 = timeDict[xx[0].attrib['TIME_SLOT_REF1']]
        time2 = timeDict[xx[0].attrib['TIME_SLOT_REF2']]
        oldAnnot.append((tName, timeDec(time1), timeDec(time2), 'x'))
        Annoz.append(int(time1))
    return Annoz

#-------------------------------------------------------------------------------
# Algorithm for finding the next 'moves' based on the current cheapest node
#-------------------------------------------------------------------------------

def boolz(n,R=[[]]):
    '''All length-n boolean lists minus the all-False list.'''
    if n < 1: return [r for r in R if any(r)]
    return boolz(n-1, [x+[1] for x in R]+[x+[0] for x in R]) 

#-------------------------------------------------------------------------------
# Dijkstra-style algorithm for finding the optimal alignment in the edit-graph
# given the assumed cost above for insertion and deletion. The only clever bit
# here is that we keep the queue sorted by (cost/totalAnnotations) in order to
# avoid exploring much beyond the diagonal of the edit-graph.  
#-------------------------------------------------------------------------------

def bestalign(Data,verbose=False):
    done = tuple([len(D)-1 for D in Data])
    zero = tuple([0 for D in Data])
    best = {zero : (0, None)}
    STATES = [ ]
    heappush(STATES, (0, 0, zero))
    (avz, cost, state) = heappop(STATES)
    #--------------------------------------------------
    # Find the cheapest edit-paths to all of the states
    #--------------------------------------------------
    while state != done:
        for incs in boolz(len(Data)):   # <-- set of possible state updates
            nuCost = cost 
            nuState = tuple([inc+oldS for (inc,oldS) in zip(incs,state)])
            if any([nuState[i] > len(Data[i]) for i in range(len(Data))]):
                continue                # <-- index too high, not a state
            #-----------------------------------------
            # If fewer than half the indices increment 
            # count all the 1's to be false positives
            #-----------------------------------------
            if sum(incs)/float(len(incs)) < 0.5:
                nuCost += skip * sum(incs)
            #-----------------------------------------
            # If half or more of the indices increment
            # count all the 0's to be missed letters
            #-----------------------------------------
            else:
                nuCost += skip * (len(incs) - sum(incs)) 
                for (idx1, S1) in enumerate(nuState):
                    for (idx2, S2) in enumerate(nuState):
                        if incs[idx1] and incs[idx2] and idx1<idx2: 
                            nuCost += abs(Data[idx1][S1] - Data[idx2][S2])
 
            nuAv = nuCost/float(sum(nuState))
            if nuState not in best or best[nuState][0] > nuCost: 
                best[nuState] = (nuCost,state)
                heappush(STATES, (nuAv, nuCost, nuState))
                # print 'Av:',int(nuAv),' cst:',nuCost,nuState, '<-- push'
        # Get the state with the best averaged cost for the next loop             
        (avz, cost, state) = heappop(STATES)
        # print 'Av:', int(avz),' cst:', cost, state,'pop <--------------------'
    # Work backward from final state using dictionary of optimal single steps 
    chain = deque([done])
    while chain[0] != zero:
        previous = best[chain[0]][-1]
        move = [a>b for (a,b) in zip(chain[0],previous)]
        if sum(move)/float(len(move)) < 0.5: 
            chain.popleft()       # Omit false positives
        else:
            times = [Data[i][pos] for (i,pos) in enumerate(chain[0]) if move[i]]
            avg = sum(times)/len(times)
            chain[0] = timeDec(avg)
        chain.appendleft(previous)
    return chain                    
            
#---------------------------------------------------------------------------
# Associate averages with letters
#-------------------------------------------------------------------------------

def makeTier(merged,spanList):
    merged.popleft()    # <- discard the zero time 
    # verbose = True
    verbose = False
    T, EndT, out = 0.0, 0.0, []
    interAv = (float(merged[-1])-float(merged[0]))/len(merged)
    while merged and spanList:          
        Av = merged.popleft()
        TT = float(Av)          
        if TT > (EndT + 0.9):
            T = TT - interAv
            S = spanList.pop(0)
            EndT = S[-1]
            SS = [sg for sg in S[1]]
            # spanStart, spanEnd = str(S[0])+'.000', str(S[2])+'.000'
            # if verbose: print 'Wd, %s, %s, %s' % (spanStart, spanEnd, S[1])
            # out.append(('Wd', spanStart, spanEnd, S[1]))
        
        if SS: Seg = SS.pop(0)
        else: Seg = '0' 
        if Seg == '(': Seg = SS.pop(0)
        Seg = Seg.upper()
        # beginS = (T + TT)/2 + 1
        # if merged: endSeg = (TT + float(merged[0]))/2 - 1
        # else: endSeg = TT + interAv/2 - 1
        beginS = str(float(Av)-0.150)
        endSeg = str(float(Av)+0.150)
        out.append(('Sg', beginS, endSeg, Seg))
        nx = str(float(Av)+0.001)
        out.append(('Av', Av, nx, Seg))
        if verbose: print 'Av, %s, %s, %s' % (Av, nx, Seg)
        T = TT
    return out


#-------------------------------------------------------------------------------
# Allows the xml output to be pretty and somewhat human readable
#-------------------------------------------------------------------------------

def indent(elem, level=0):
    i = "\n" + level*"    "
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + "    "
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
        for elem in elem:
            indent(elem, level+1)
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
    else:
        if level and (not elem.tail or not elem.tail.strip()):
            elem.tail = i


#-------------------------------------------------------------------------------
# Build an xml file that elan likes
#-------------------------------------------------------------------------------

def buildEAF(nuTiers):
    tree = ET.parse(sys.argv[2])
    root = tree.getroot()

    # #----------------------------------------------------
    # # Build the tree and fill in the basic attributes
    # #----------------------------------------------------            
    # root = ET.Element("ANNOTATION_DOCUMENT", AUTHOR="", DATE="2010-04-16T10:47:56-06:00", FORMAT="2.6", VERSION="2.6")
    # root.attrib["xmlns:xsi"] = "http://www.w3.org/2001/XMLSchema-instance"
    # root.attrib["xsi:noNamespaceSchemaLocation"] = "http://www.mpi.nl/tools/elan/EAFv2.6.xsd"
    # 
    # header = ET.SubElement(root, "HEADER", MEDIA_FILE="", TIME_UNITS="milliseconds")
    # 
    # media_desc = ET.SubElement(header, "MEDIA_DESCRIPTOR", MEDIA_URL="file:///Users/aslr/Desktop/Andy100307Careful_HF10_Segment1.mov", MIME_TYPE="video/quicktime", RELATIVE_MEDIA_URL="file:/../../../Andy100307Careful_HF10_Segment1.mov")
    # 
    # prop = ET.SubElement(header, "PROPERTY", NAME="lastUsedAnnotationId")
    # prop.text ="642"
    # 
    # ## We might need to set this <--------------------------------
    # 
    # time_order = ET.SubElement(root, "TIME_ORDER")
    # 
    # linguistic_type = ET.SubElement(root, "LINGUISTIC_TYPE", GRAPHIC_REFERENCES="false", LINGUISTIC_TYPE_ID="default-lt", TIME_ALIGNABLE="true")
    # 
    # locale = ET.SubElement(root, "LOCALE", COUNTRY_CODE="US", LANGUAGE_CODE="en")
    # 
    # constraint = ET.SubElement(root, "CONSTRAINT", DESCRIPTION="Time subdivision of parent annotation's time interval, no time gaps allowed within this interval", STEREOTYPE="Time_Subdivision")
    # 
    # constraint = ET.SubElement(root, "CONSTRAINT", DESCRIPTION="Symbolic subdivision of a parent annotation. Annotations refering to the same parent are ordered", STEREOTYPE="Symbolic_Subdivision")
    # 
    # constraint = ET.SubElement(root, "CONSTRAINT", DESCRIPTION="1-1 association with a parent annotation", STEREOTYPE="Symbolic_Association")
    # 
    # constraint = ET.SubElement(root, "CONSTRAINT", DESCRIPTION="Time alignable annotations within the parent annotation's time interval, gaps are allowed", STEREOTYPE="Included_In" )

    root.remove(root[2])
    time_order = root[1]
    for oldTO in time_order.getchildren():
        time_order.remove(oldTO)

    #----------------------------------------------------
    # build the tiers from the contents of nuTiers
    #----------------------------------------------------            
    knownTiers = {}
    tslt = 0
    anot = 0
    
    for Name, M, nx, Seg in nuTiers:
        # if Name == 'Wd': continue
        tslt += 1
        time_slot_id0 = 'ts' + Name + str(tslt) 
        M = int(float(M)*1000)
        time_value0 = str(M)
        time_slot = ET.SubElement(time_order, "TIME_SLOT")
        time_slot.attrib["TIME_SLOT_ID"] = time_slot_id0
        time_slot.attrib["TIME_VALUE"] = time_value0
    
        tslt += 1
        time_slot_id1 = 'ts' + Name + str(tslt) 
        nx = int(float(nx)*1000)
        time_value1 = str(nx)
        time_slot = ET.SubElement(time_order, "TIME_SLOT")
        time_slot.attrib["TIME_SLOT_ID"] = time_slot_id1
        time_slot.attrib["TIME_VALUE"] = time_value1
    
        if Name not in knownTiers:
            #----------------------------------------------------
            # Create a tier
            #----------------------------------------------------            
            tier = ET.SubElement(root, "TIER")
            tier.attrib["ANNOTATOR"] = 'transMerge'
            tier.attrib["DEFAUTL_LOCALE"] = "en"
            tier.attrib["LINGUISTIC_TYPE_REF"] = "default-lt"
            tier.attrib["TIER_ID"] = Name
            
            annotation = ET.SubElement(tier, "ANNOTATION")
            
            knownTiers[Name] = (tier, annotation)
            #----------------------------------------------------    
                    
        tier, annotation = knownTiers[Name]      
        anot += 1
        annotation_id = 'a' + str(anot)
        alignable_annotation = ET.SubElement(annotation, "ALIGNABLE_ANNOTATION")
        alignable_annotation.attrib["ANNOTATION_ID"] = annotation_id
        alignable_annotation.attrib["TIME_SLOT_REF1"] = time_slot_id0
        alignable_annotation.attrib["TIME_SLOT_REF2"] = time_slot_id1
    
        anno_value = ET.SubElement(alignable_annotation, "ANNOTATION_VALUE")
        anno_value.text = Seg
        
    #----------------------------------------------------
    # Output a new .eaf file
    #----------------------------------------------------            
    newFileName = 'ALG' + str(sys.argv[2])[2:]
    indent(root)
    tree = ET.ElementTree(root) 
    tree.write(newFileName, encoding="utf-8")
    print ".eaf file written: ", newFileName
    

#-------------------------------------------------------------------------------
# Launch the main program loop
#-------------------------------------------------------------------------------

main()